﻿#' @title 单多因素COX回归分析，R包绘制 三线表
#' @param clinical 临床信息文件，需要time 与status两列作为时间信息
#' @param SaveFile 是否在输出目录保存图片
#' @param ad 结果输出目录，生成figure与table文件夹，分别存储图片与表格文件,字符串向量，必须以 **"/"** 结尾
#' @param var_name 用于命名文件夹，如果为NULL，则使用clinical变量名
#' @param w 图宽
#' @param h 图高
#' @param p_cut 挑选符合p之的单因素元素进行多因素
#' @param features_col 字符串向量,纳入分析的特征名称，必须为`clinical`中的列名
#' @param cohort 队列信息，默认是TCGA
#' @export 
#' @usage 
#' a <- Forest_DF(
#'    clinical = forest_data, features_col = c("Score","HRD_Score", "Age", "Stage", "Menopause_status"),
#'    SaveFile = T, od = file.path(out_home, "out/3.lasso_model/forest_df/"), w = 5,
#'    h = 2.5, p_cut = 2, var_name = 'tnbc_tcga'
#' )
#' 
#' @return  *list*
#' @author WYK
#' 
Forest_DF <- ClinicalCox_DF_v3 <- function(
    clinical = NULL,
    features_col = c("Smoking", "EGFR", "ALK_eml4", "KRAS", "Histological.Type", "Score"),
    SaveFile = T, od = "./", w = 5, h = 3, p_cut = 2, var_name = NULL, cohort = "TCGA") {
    require(tidyverse)
    require(survival)
    require(survminer)
    require(cowplot)

    if (is.null(var_name)) {
        var_name <- Sys.time() %>% as.character()
    }

    levels <- map(features_col, \(x){
        clinical %>%
            select(x) %>%
            pull() %>%
            factor() %>%
            as.list() %>%
            lvls_union()
    })
    names(levels) <- features_col

    features_chr <- map(features_col, \(x){
        # x <- 'TCGA_Subtype'
        str_c(x, " ", paste0("(", levels[[x]][-1], " vs. ", levels[[x]][1], ")"))
    }) %>%
        unlist() %>%
        as_tibble()

    vars_for_table <- clinical %>%
        dplyr::select(any_of(features_col)) %>%
        colnames(.)

    univ_formulas <- sapply(vars_for_table, function(x) as.formula(paste("Surv(time, status)~", x)))
    univ_models <- lapply(univ_formulas, function(x) {
        coxph(x, data = clinical)
    })

    univcox_clinical_res <- map_dfr(.x = univ_models, .f = function(x) {
        result <- summary(x)

        result2 <- data.frame(
            var = rownames(result$coefficients),
            coef = result$coefficients[, 1],
            Hazard_Ratio = result$coefficients[, 2],
            se = result$coefficients[, 3],
            # p.value = result$coefficients[, 5],
            p.value = result$sctest["pvalue"],
            lower_.95 = result$conf.int[, "lower .95"],
            upper_.95 = result$conf.int[, "upper .95"],
            logrank_pvalue = result$sctest["pvalue"],
            wald_pvalue = result$waldtest["pvalue"],
            Likelihood_pvalue = result$logtest["pvalue"]
        )

        result2 <- result2 %>%
            mutate(HR = paste0(
                sprintf("%.2f", Hazard_Ratio),
                "(",
                sprintf("%.2f", lower_.95),
                "-",
                sprintf("%.2f", upper_.95),
                ")"
            ))
    })

    rownames(univcox_clinical_res) <- univcox_clinical_res$var
    univcox_clinical_res$var <- features_chr$value

    univ_cox_df <- univcox_clinical_res %>%
        mutate(P = ifelse(p.value < 0.001, "<0.001", sprintf("%.3f", p.value))) %>%
        select(var, HR, P) %>%
        dplyr::rename("HR(95%CI)" = HR)

    p_univ_cox <- univ_cox_df %>%
        dplyr::rename(`Univariable Cox` = var) %>%
        ggtexttable(., rows = NULL, theme = ttheme("blank")) %>%
        tab_add_hline(at.row = 1:2, row.side = "top", linewidth = 2) %>%
        tab_add_hline(at.row = dim(univ_cox_df)[1] + 1, row.side = "bottom", linewidth = 2) %>%
        tab_add_title(
            text = cohort,
            face = "bold",
            size = 16,
            padding = unit(1.5, "line"),
            just = "left"
        ) +
        theme(text = element_text(family = "Times New Roman"))
    # Univariable Cox
    # Univariate

    # MultiCox Part-------------

    # p_cut <- .05
    do_feature_select <- function(p_cut) {
        col_names_left <- map(univ_cox_df %>% filter(P < p_cut) %>% pull(1) %>% str_split(" \\("), ~ .x[[1]]) %>%
            unlist() %>%
            unique()

        vars_for_table <- clinical %>%
            dplyr::select(any_of(col_names_left)) %>%
            colnames(.)

        levels <- map(col_names_left, \(x){
            clinical %>%
                select(x) %>%
                pull() %>%
                factor() %>%
                as.list() %>%
                lvls_union()
        })
        names(levels) <- col_names_left

        features_chr <- map(col_names_left, \(x){
            # x <- 'TCGA_Subtype'
            str_c(x, " ", paste0("(", levels[[x]][-1], " vs. ", levels[[x]][1], ")"))
        }) %>%
            unlist() %>%
            as_tibble()

        return(list("features_chr" = features_chr, "col_names_left" = col_names_left))
    }

    features_new <- do_feature_select(p_cut)

    if (nrow(features_new$features_chr) < 2) {
        features_chr <- features_chr
        vars_for_table <- vars_for_table

        warning(str_glue("在p小于 {crayon::blue(p_cut)} 的时候，只有一个临床信息保留，所以使用全部变量进行多因素Cox分析。"))
    } else {
        features_chr <- features_new$features_chr
        vars_for_table <- features_new$col_names_left
    }


    multicox_formulas <-
        as.formula(paste(
            "Surv(time, status)~",
            paste0(sep = "`", vars_for_table, sep = "`", collapse = "+")
        ))

    multicox_cox_res <- coxph(
        formula = multicox_formulas,
        data = clinical
    )

    multicox_cox_res_summary <- summary(multicox_cox_res)

    multicox_cox_res_df <- data.frame(
        var = rownames(multicox_cox_res_summary$coefficients),
        coef = multicox_cox_res_summary$coefficients[, 1],
        Hazard_Ratio = multicox_cox_res_summary$coefficients[, 2],
        se = multicox_cox_res_summary$coefficients[, 3],
        p.value = multicox_cox_res_summary$coefficients[, 5],
        lower_.95 = multicox_cox_res_summary$conf.int[, "lower .95"],
        upper_.95 = multicox_cox_res_summary$conf.int[, "upper .95"],
        logrank_pvalue = multicox_cox_res_summary$sctest["pvalue"],
        wald_pvalue = multicox_cox_res_summary$waldtest["pvalue"],
        Likelihood_pvalue = multicox_cox_res_summary$logtest["pvalue"]
    )

    multicox_cox_res_df <- multicox_cox_res_df %>%
        mutate(HR = paste0(
            sprintf("%.2f", Hazard_Ratio),
            "(",
            sprintf("%.2f", lower_.95),
            "-",
            sprintf("%.2f", upper_.95),
            ")"
        ))

    multicox_cox_res_df$var <- features_chr$value

    multi_cox_df <- multicox_cox_res_df %>%
        mutate(P = ifelse(p.value < 0.001, "<0.001", sprintf("%.3f", p.value))) %>%
        select(var, HR, P) %>%
        dplyr::rename("HR(95%CI)" = HR)

    p_multi_cox <- multi_cox_df %>%
        dplyr::rename(`Multivariable Cox` = var) %>%
        ggtexttable(., rows = NULL, theme = ttheme("blank")) %>%
        tab_add_hline(at.row = 1:2, row.side = "top", linewidth = 2) %>%
        tab_add_hline(at.row = dim(multi_cox_df)[1] + 1, row.side = "bottom", linewidth = 2) %>%
        tab_add_title(
            text = cohort,
            face = "bold",
            size = 16,
            padding = unit(1.5, "line"),
            just = "left"
        ) +
        theme(text = element_text(family = "Times New Roman"))
    # Multivariate Cox

    p_all <- ggpubr::ggarrange(p_univ_cox, p_multi_cox, nrow = 2, ncol = 1)

    # title <- ggdraw() + draw_label("TCGA", fontface='bold' ,hjust = 0, vjust = 0) #nrow(univ_cox_df)
    # ggarrange(title,p_all,ncol = 1,height = c(.1,1))


    if (isTRUE(SaveFile)) {
        if (!dir.exists(od)) {
            dir.create(od, recursive = T)
        }

        walk(c("p_univ_cox", "p_multi_cox"), function(x) {
            ggsave(plot = get0(x), filename = str_glue("{od}/Figure_{x}_{var_name}.pdf"), width = w, height = h)
        })

        ggsave(plot = p_all, filename = str_glue("{od}/Figure_Univ_Multi_Cox_{var_name}.pdf"), width = w, height = h * 1.85)

        walk(c("univ_cox_df", "univcox_clinical_res", "multicox_cox_res_df", "multi_cox_df"), function(x) {
            write.table(x = get0(x), file = str_glue("{od}/Table_{x}_{var_name}.txt"), quote = F, sep = "\t")
        })
    }

    return(list("univcox_res" = univcox_clinical_res, "multicox_res" = multicox_cox_res_df))
}

# setdiff(colnames(aa), c("sample", "time", "status", "riskscore")) %>%
#     str_c('"', ., '"') %>%
#     paste0(",") %>%
#     cat()


# a <- Forest_DF(
#    clinical = forest_data, features_col = c("Score","HRD_Score", "Age", "Stage", "Menopause_status"),
#    SaveFile = F, od = file.path(out_home, "out/3.lasso_model/forest_df/"), w = 5,
#    h = 2.5, p_cut = 2, var_name = 'tnbc_tcga'
# )